package com.olink;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.security.MessageDigest;
import java.util.*;

import java.util.Base64;

public class WebSocketServer {

    private Selector selector;
    private ServerSocketChannel serverSocketChannel;
    private static final int PORT = 8080;

    public static void main(String[] args) throws IOException {
        new WebSocketServer().start();
    }

    public void start() throws IOException {
        selector = Selector.open();
        serverSocketChannel = ServerSocketChannel.open();
        serverSocketChannel.bind(new InetSocketAddress(PORT));
        serverSocketChannel.configureBlocking(false);
        serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("WebSocket NIO 服务器启动，监听端口 " + PORT);

        while (true) {
            selector.select(); // 阻塞直到有事件
            Set<SelectionKey> keys = selector.selectedKeys();
            Iterator<SelectionKey> iterator = keys.iterator();

            while (iterator.hasNext()) {
                SelectionKey key = iterator.next();
                iterator.remove(); // 移除以避免重复处理

                try {
                    if (key.isAcceptable()) {
                        handleAccept(key);
                    } else if (key.isReadable()) {
                        handleRead(key);
                    }
                } catch (Exception e) {
                    key.cancel();
                    key.channel().close();
                    System.err.println("连接处理异常：" + e.getMessage());
                }
            }
        }
    }

    // 处理客户端连接
    private void handleAccept(SelectionKey key) throws IOException {
        ServerSocketChannel ssc = (ServerSocketChannel) key.channel();
        SocketChannel client = ssc.accept();
        client.configureBlocking(false);
        client.register(selector, SelectionKey.OP_READ, new ClientAttachment());
        System.out.println("新客户端连接：" + client.getRemoteAddress());
    }

    // 处理客户端数据
    private void handleRead(SelectionKey key) throws IOException {
        SocketChannel client = (SocketChannel) key.channel();
        ClientAttachment attachment = (ClientAttachment) key.attachment();
        ByteBuffer buffer = ByteBuffer.allocate(1024);
        int read = client.read(buffer);

        if (read == -1) {
            client.close();
            System.out.println("客户端关闭连接");
            return;
        }

        buffer.flip();
        byte[] data = new byte[buffer.remaining()];
        buffer.get(data);

        if (!attachment.handshakeDone) {
            String request = new String(data);
            if (request.contains("Sec-WebSocket-Key")) {
                String response = buildHandshakeResponse(request);
                client.write(ByteBuffer.wrap(response.getBytes()));
                attachment.handshakeDone = true;
                System.out.println("握手完成：" + client.getRemoteAddress());
            }
        } else {
            // 处理 WebSocket 数据帧
            String message = decodeWebSocketFrame(data);

            System.out.println("收到客户端消息: " + message);


            byte[] response = encodeWebSocketFrame("客户端"+ client.getRemoteAddress()+"发送了"+message);

            client.write(ByteBuffer.wrap(response));
        }
    }

    // 构造 WebSocket 握手响应
    private String buildHandshakeResponse(String request) {
        String keyLine = Arrays.stream(request.split("\r\n"))
                .filter(line -> line.startsWith("Sec-WebSocket-Key:"))
                .findFirst()
                .orElse("");
        String key = keyLine.split(":")[1].trim();
        String accept = generateAcceptKey(key);

        return "HTTP/1.1 101 Switching Protocols\r\n" +
                "Upgrade: websocket\r\n" +
                "Connection: Upgrade\r\n" +
                "Sec-WebSocket-Accept: " + accept + "\r\n\r\n";
    }

    // 生成 Sec-WebSocket-Accept
    private String generateAcceptKey(String key) {
        try {
            String magicString = key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
            MessageDigest digest = MessageDigest.getInstance("SHA-1");
            byte[] hash = digest.digest(magicString.getBytes("UTF-8"));
            return Base64.getEncoder().encodeToString(hash);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // 解码 WebSocket 文本帧（只支持文本帧）
    private String decodeWebSocketFrame(byte[] data) {
        int payloadLen = data[1] & 0x7F;
        int maskIndex = 2;
        if (payloadLen == 126) maskIndex = 4;
        else if (payloadLen == 127) maskIndex = 10;

        byte[] mask = Arrays.copyOfRange(data, maskIndex, maskIndex + 4);
        int payloadStart = maskIndex + 4;
        byte[] payload = Arrays.copyOfRange(data, payloadStart, data.length);

        for (int i = 0; i < payload.length; i++) {
            payload[i] = (byte) (payload[i] ^ mask[i % 4]);
        }
        return new String(payload);
    }


    private byte[] encodeWebSocketFrame(String message) {
        byte[] data = message.getBytes();
        int maxFrameSize = 60; // 每帧最多60字节
        List<byte[]> frames = new ArrayList<>();

        int pos = 0;
        while (pos < data.length) {
            int len = Math.min(maxFrameSize, data.length - pos);
            byte[] subPayload = Arrays.copyOfRange(data, pos, pos + len);

            boolean isFirst = (pos == 0);
            boolean isFinal = (pos + len >= data.length);

            byte firstByte;
            if (isFirst && isFinal) {
                firstByte = (byte) 0x81; // FIN=1，文本帧
            } else if (isFirst) {
                firstByte = 0x01;        // FIN=0，文本帧（首帧）
            } else if (isFinal) {
                firstByte = (byte) 0x80; // FIN=1，继续帧（末帧）
            } else {
                firstByte = 0x00;        // FIN=0，继续帧（中间帧）
            }

            ByteBuffer frame;

            if (subPayload.length <= 125) {
                // 长度 <= 125，长度字段1字节
                frame = ByteBuffer.allocate(2 + subPayload.length);
                frame.put(firstByte);
                frame.put((byte) subPayload.length);
                frame.put(subPayload);
            } else if (subPayload.length <= 0xFFFF) {
                // 长度 <= 65535，长度字段用126 + 2字节扩展长度
                frame = ByteBuffer.allocate(4 + subPayload.length);
                frame.put(firstByte);
                frame.put((byte) 126);
                // 2字节无符号短整型大端序存长度
                frame.putShort((short) subPayload.length);
                frame.put(subPayload);
            } else {
                // 长度 > 65535，长度字段用127 + 8字节扩展长度
                frame = ByteBuffer.allocate(10 + subPayload.length);
                frame.put(firstByte);
                frame.put((byte) 127);
                // 8字节无符号长整型大端序存长度
                frame.putLong(subPayload.length);
                frame.put(subPayload);
            }

            frames.add(frame.array());
            pos += len;
        }

        // 合并所有帧
        int totalLength = frames.stream().mapToInt(f -> f.length).sum();
        ByteBuffer result = ByteBuffer.allocate(totalLength);
        for (byte[] f : frames) {
            result.put(f);
        }

        return result.array();
    }



    // 每个客户端附加对象
    static class ClientAttachment {
        boolean handshakeDone = false;
    }
}
